#include"ORCL.h"

struct orcl_context* orcl_context_init(int node_id, int group_id, int node_num, int group_num, int group_size, int dci_num, int port_index, 
int numa_node_id, volatile void *prealloc_mes_buf, int mes_buf_size, int shm_key_start, volatile orcl_context* pre_ctx, int orcl_ctx_id){
    int node_cpu_tp[NODE_COUNT][NODE_CPU_COUNT] = {{0,1,2,3,4,5,12,13,14,15,16,17}, {6,7,8,9,10,11,18,19,20,21,22,23}};

    assert(port_index >= 0 && port_index <= 16);
	assert(numa_node_id >= -1 && numa_node_id <= 8);
	assert(dci_num >= 0 && dci_num <= group_num);
	assert(mes_buf_size >= 0 && mes_buf_size <= M_1024);
    assert(node_id >= 0 && node_id < node_num);
    assert(group_num >= 0 && group_num <= node_num);
    assert(group_id >= 0 && group_id < group_num);
    
    // 如果之前当前节点没有申请过ctx，pre_ctx==NULL
    struct orcl_context* orcl_ctx = (struct orcl_context *)malloc(sizeof(struct orcl_context));
	memset(orcl_ctx, 0, sizeof(struct orcl_context));

    // 三种发送计数器初始化，计数器是独占的
    memset(orcl_ctx->rc_forward_signal_period, 0, ORCL_MAX_GROUP_SIZE * sizeof(int));
    memset(orcl_ctx->rc_normal_signal_period, 0, ORCL_MAX_GROUP_SIZE * sizeof(int));
    memset(orcl_ctx->dci_signal_period, 0, ORCL_MAX_DCI_NUM * sizeof(int));

    orcl_ctx->node_id = node_id;
    orcl_ctx->group_id = group_id;
    orcl_ctx->node_num = node_num;
    orcl_ctx->group_num = group_num;
    orcl_ctx->dci_num = dci_num;
    orcl_ctx->message_buffer_size = mes_buf_size;
    orcl_ctx->shm_key_start = shm_key_start;
    // shm_key_end代表下一个要使用的shm_key
    orcl_ctx->shm_key_end = shm_key_start;
    orcl_ctx->port_index = port_index;
    orcl_ctx->group_size = group_size;
    orcl_ctx->numa_node_id = numa_node_id;
    orcl_ctx->orcl_ctx_id = orcl_ctx_id;

    // 开启可以开启的所有权限
    int ib_flags = IBV_ACCESS_LOCAL_WRITE | IBV_ACCESS_REMOTE_READ | 
        IBV_ACCESS_REMOTE_WRITE | IBV_ACCESS_REMOTE_ATOMIC;


    // 在最外面申请消息缓冲区
    if(prealloc_mes_buf == NULL){
        int reg_size = 0;
                
        // 为了使用大页，使用共享内存
        while(reg_size < mes_buf_size){
            reg_size += M_2;
        }

        assert(orcl_ctx->shm_key_end >= 1 && orcl_ctx->shm_key_end <= 128);

        // 对于用户不透明的message，在申请normal的连接的时候就申请好
        orcl_ctx->message_buffer = (volatile uint8_t *) orcl_malloc_socket(
            orcl_ctx->shm_key_end, reg_size, numa_node_id);

        orcl_ctx->shm_key_end = orcl_ctx->shm_key_end + 1;
        
        prealloc_mes_buf = orcl_ctx->message_buffer;

        // 初始化
        memset((char *) orcl_ctx->message_buffer, 0, reg_size);
    }else{
        orcl_ctx->message_buffer = (volatile uint8_t *)prealloc_mes_buf;
    }


    int i;

    // 首先申请group_num个转发缓冲区
    // 这个转发缓冲区是私有的，每个转发线程负责一个
    // 转发线程数量和组数量相同
    for(i = 0; i < group_num; i++){
        // 当前组不需要自己对自己的转发
        if(group_id == i){
            continue;
        }

        orcl_ctx->rc_forward_buf[i] = (volatile uint8_t *) orcl_malloc_socket(
            orcl_ctx->shm_key_end, FORWARD_RC_BUF_SIZE, numa_node_id);
        
        orcl_ctx->shm_key_end = orcl_ctx->shm_key_end + 1;

        // 初始化
        memset((char *) orcl_ctx->rc_forward_buf[i], 0, FORWARD_RC_BUF_SIZE);
    }


    if(pre_ctx == NULL){
        // 首先创建所有RC相关的上下文，一共有group_size个rc，中间有一个不用实际初始化
        // 一个节点在当前组中的编号是node_id%group_size
        // 组内索引为i的节点的node_id是group_size*group_id+i
        // 都是0base的计数方式
        // 首先处理所有的normal_rc，这些rc唯一的缓冲区就是对于用户透明的message_buf
        for(i = 0; i < group_size; i++){
            if(node_id % group_size == i){
                // 跳过与自己的连接
                continue;
            }
            
            // 不是自己，那么创造自己的上下文，并且发布自己的队列信息
            // 首先找设备
            struct ibv_device *ib_dev = orcl_resolve_port_index(orcl_ctx, port_index);
            printf("find device %s\n", ibv_get_device_name(ib_dev));
            CPE(!ib_dev, "ORCL: IB device not found", 0);

            // 打开设备，创造上下文，并且记录下来
            orcl_ctx->rc_normal_ctx[i] = ibv_open_device(ib_dev);
            CPE(!orcl_ctx->rc_normal_ctx[i], "ORCL: Couldn't get context", 0);

            // 为这个rc创建保护域
            orcl_ctx->rc_normal_pd[i] = ibv_alloc_pd(orcl_ctx->rc_normal_ctx[i]);
            CPE(!orcl_ctx->rc_normal_pd[i], "ORCL: Couldn't allocate PD", 0);

            // 创建对应的rc队列和cq
            orcl_create_normal_conn_qps(orcl_ctx, i);

            // 之前已经申请了message缓冲区，这里进行注册
            orcl_ctx->rc_normal_mes_mr[i] = ibv_reg_mr(orcl_ctx->rc_normal_pd[i], (char *) orcl_ctx->message_buffer, orcl_ctx->message_buffer_size, ib_flags);
            assert(orcl_ctx->rc_normal_mes_mr[i] != NULL);
            
        }

        printf("init: init normal rc qp node_id:%d, ctx_id:%d\n", node_id, orcl_ctx_id);

        // 接下来初始化所有的rc_forward内容，初始化的方式和normal基本上就是一样的
        // 与normal节点不同，forward节点涉及缓冲区
        for(i = 0; i < group_size; i++){
            // 跳过自己
            if(node_id % group_size == i){
                // 跳过与自己的连接
                continue;
            }
            
            // 不是自己，那么创造自己的上下文，并且发布自己的队列信息
            // 首先找设备
            struct ibv_device *ib_dev = orcl_resolve_port_index(orcl_ctx, port_index);
            printf("find device %s\n", ibv_get_device_name(ib_dev));
            CPE(!ib_dev, "ORCL: IB device not found", 0);

            // 打开设备，创造上下文，并且记录下来
            orcl_ctx->rc_forward_ctx[i] = ibv_open_device(ib_dev);
            CPE(!orcl_ctx->rc_forward_ctx[i], "ORCL: Couldn't get context", 0);

            // 为这个rc创建保护域
            orcl_ctx->rc_forward_pd[i] = ibv_alloc_pd(orcl_ctx->rc_forward_ctx[i]);
            CPE(!orcl_ctx->rc_forward_pd[i], "ORCL: Couldn't allocate PD", 0);

            // 开启可以开启的所有权限
            int ib_flags = IBV_ACCESS_LOCAL_WRITE | IBV_ACCESS_REMOTE_READ | 
		        IBV_ACCESS_REMOTE_WRITE | IBV_ACCESS_REMOTE_ATOMIC;

            // 创建forward队列
            orcl_create_forward_conn_qps(orcl_ctx, i);

            // 下面首先注册message缓冲区，在这里message缓冲区肯定是已经被申请了，只需要执行注册就可以
            orcl_ctx->rc_forward_mes_mr[i] = ibv_reg_mr(orcl_ctx->rc_forward_pd[i], (char *) orcl_ctx->message_buffer, orcl_ctx->message_buffer_size, ib_flags);
			assert(orcl_ctx->rc_forward_mes_mr[i] != NULL);            
            
            // rc_forward_mr一共group num组，每组group size个
            // 我们要初始化的每一组的第i个
            // 转发缓冲区要注册到所有forward rc，作为forward rc的发送buffer
            int j;
            for(j = 0; j < group_num; j++){
                // 同一组的不需要转发
                if(j == group_id){
                    // 跳过
                    continue;
                }

                // 转发RC数量和group size相同，转发缓冲区数量

                errno = 0;
                // 注册
                orcl_ctx->rc_forward_mr[j][i] = ibv_reg_mr(orcl_ctx->rc_forward_pd[i], (char *) orcl_ctx->rc_forward_buf[j], FORWARD_RC_BUF_SIZE, ib_flags);

                if(errno == EINVAL){
                    printf("Invalid access value\n");
                }else if(errno == ENOMEM){
                    printf("Not enough resources (either in operating system or in RDMA device) to complete this operation\n");
                }


                assert(orcl_ctx->rc_forward_mr[j][i] != NULL);
            }			
        }

        printf("init: init forward rc qp node_id:%d, ctx_id:%d\n", node_id, orcl_ctx_id);

        orcl_publish_conn_qp(orcl_ctx);

        printf("init: publish all rc qp info node_id:%d, ctx_id:%d\n", node_id, orcl_ctx_id);
    }else{
        // 复用之前的rc相关内容，包含forward_rc和normal_rc
        int i;
        for(i = 0; i < group_size; i++){
            // 跳过自己
            if(orcl_ctx->node_id % orcl_ctx->group_size == i){
                // 跳过与自己的连接
                continue;
            }

            // 这里查看复制过去的东西有没有问题
            assert(pre_ctx->rc_normal_ctx[i] != NULL && pre_ctx->rc_normal_pd[i] != NULL
                && pre_ctx->rc_normal_qp[i] != NULL);

            // 设备id和端口id也是共享的
            orcl_ctx->device_id = pre_ctx->device_id;
            orcl_ctx->dev_port_id = pre_ctx->dev_port_id;

            // 只有qp之间相关的元数据是共享的
            orcl_ctx->rc_normal_ctx[i] = pre_ctx->rc_normal_ctx[i];
            orcl_ctx->rc_normal_pd[i] = pre_ctx->rc_normal_pd[i];
            orcl_ctx->message_buffer = pre_ctx->message_buffer;
            orcl_ctx->message_buffer_size = pre_ctx->message_buffer_size;
            orcl_ctx->rc_normal_mes_mr[i] = pre_ctx->rc_normal_mes_mr[i];
            orcl_ctx->rc_normal_qp[i] = pre_ctx->rc_normal_qp[i];
            orcl_ctx->rc_normal_cq[i] = pre_ctx->rc_normal_cq[i];

            orcl_ctx->rc_forward_ctx[i] = pre_ctx->rc_forward_ctx[i];
            orcl_ctx->rc_forward_pd[i] = pre_ctx->rc_forward_pd[i];
            orcl_ctx->rc_forward_mes_mr[i] = pre_ctx->rc_forward_mes_mr[i];
            orcl_ctx->rc_forward_qp[i] = pre_ctx->rc_forward_qp[i];
            orcl_ctx->rc_forward_cq[i] = pre_ctx->rc_forward_cq[i];

            // forward buf和注册等内容不是共享的
            // 转发缓冲区要注册到所有forward rc，作为forward rc的发送buffer
            int j;
            for(j = 0; j < group_num; j++){
                // 同一组的不需要转发
                if(j == group_id){
                    // 跳过
                    continue;
                }

                // 查看对应值对不对
                assert(orcl_ctx->rc_forward_pd[i] != NULL);

                // 注册
                orcl_ctx->rc_forward_mr[j][i] = ibv_reg_mr(orcl_ctx->rc_forward_pd[i], (char *) orcl_ctx->rc_forward_buf[j], FORWARD_RC_BUF_SIZE, ib_flags);
                assert(orcl_ctx->rc_forward_mr[j][i] != NULL);
            }

            orcl_publish_conn_qp(orcl_ctx);
            
            printf("init: publish all rc qp info node_id:%d, ctx_id:%d\n", node_id, orcl_ctx_id);
        }

        
    }

    // dci永远是独享的
    // 首先创建dci num个dci
    orcl_ctx->dci_num = dci_num;
    orcl_ctx->next_dci_id = 0;
    
    for(i = 0; i < dci_num; i++){
        // 首先创建对应上下文
        struct ibv_device *ib_dev = orcl_resolve_port_index(orcl_ctx, port_index);
	    CPE(!ib_dev, "ORCL: IB device not found", 0);

        orcl_ctx->dci_ctx[i] = ibv_open_device(ib_dev);
	    CPE(!orcl_ctx->dci_ctx[i], "ORCL: Couldn't get context", 0);

        // 申请保护域
        orcl_ctx->dci_pd[i] = ibv_alloc_pd(orcl_ctx->dci_ctx[i]);
	    CPE(!orcl_ctx->dci_pd[i], "ORCL: Couldn't allocate PD", 0);

        // 将message buf注册给dci，message buf是唯一的dci发送缓冲区，除了这个之外，没有更多用途。
        orcl_ctx->dci_mes_mr[i] = ibv_reg_mr(orcl_ctx->dci_pd[i], (char *) orcl_ctx->message_buffer, orcl_ctx->message_buffer_size, ib_flags);

        // 创建dci
        orcl_create_dc_initiators(orcl_ctx, i);
    }

    printf("init: create all dci node_id:%d, ctx_id:%d\n", node_id, orcl_ctx_id);

    // 初始化两个dci对应的dm
    for(i = 0; i < dci_num; i++){
        // 申请dm
        struct ibv_exp_alloc_dm_attr dm_attr;
        memset(&dm_attr, 0, sizeof(struct ibv_exp_alloc_dm_attr));
        dm_attr.length = DCI_DM_SIZE;
        orcl_ctx->dci_dm[i] = ibv_exp_alloc_dm(orcl_ctx->dci_ctx[i], &dm_attr);
        assert(orcl_ctx->dci_dm[i] != NULL);

        // 注册
        struct ibv_exp_reg_mr_in mr_in;
        memset(&mr_in, 0, sizeof(struct ibv_exp_reg_mr_in));
        mr_in.pd = orcl_ctx->dci_pd[i];
        mr_in.addr = 0;
        mr_in.length = DCI_DM_SIZE;
        mr_in.exp_access = IBV_EXP_ACCESS_LOCAL_WRITE;
        mr_in.create_flags = 0;
        mr_in.dm = orcl_ctx->dci_dm[i];
        mr_in.comp_mask = IBV_EXP_REG_MR_DM;
        orcl_ctx->dci_dm_mr[i] = ibv_exp_reg_mr(&mr_in);
        assert(orcl_ctx->dci_dm_mr[i] != NULL);
    }

    printf("init: create all dci dm node_id:%d, ctx_id:%d\n", node_id, orcl_ctx_id);

    // 创建一系列DCT，当前组不需要dct
    for(i = 0; i < group_num; i++){
        if(orcl_ctx->group_id == i){
            continue;
        }

        // 首先申请DCT上下文
        struct ibv_device *ib_dev = orcl_resolve_port_index(orcl_ctx, port_index);
	    CPE(!ib_dev, "ORCL: IB device not found", 0);

        orcl_ctx->dct_ctx[i] = ibv_open_device(ib_dev);
	    CPE(!orcl_ctx->dct_ctx[i], "ORCL: Couldn't get context", 0);

        orcl_ctx->dct_pd[i] = ibv_alloc_pd(orcl_ctx->dct_ctx[i]);
	    CPE(!orcl_ctx->dct_pd[i], "ORCL: Couldn't allocate PD", 0);

        // 创建转发机构的接收区
        orcl_ctx->dct_forward_buf[i] = (volatile uint8_t *) orcl_malloc_socket(
            orcl_ctx->shm_key_end, FORWARD_DCT_BUF_SIZE, numa_node_id);
        
        orcl_ctx->shm_key_end = orcl_ctx->shm_key_end + 1;

        // 初始化
        memset((char *) orcl_ctx->dct_forward_buf[i], 0, FORWARD_DCT_BUF_SIZE);

        // 将转发缓冲区注册给dct
        orcl_ctx->dct_forward_mr[i] = ibv_reg_mr(orcl_ctx->dct_pd[i], (char *) orcl_ctx->dct_forward_buf[i], FORWARD_DCT_BUF_SIZE, ib_flags);

        // 创建dct target
        orcl_create_dc_target(orcl_ctx, i);
    }

    // 创建完之后发布dct的基本信息
    orcl_publish_dc_target(orcl_ctx);

    // 将两个方向的路由表初始化，一开始全部初始化为-1
    // 首先是每个dci与group的对应关系
    for(i = 0; i < dci_num; i++){
        orcl_ctx->dci_connect_group_id[i] = -1;
    }

    for(i = 0; i < group_num; i++){
        orcl_ctx->group_connect_dci_id[i] = -1;
    }

    printf("init: publish all dc target node_id:%d, ctx_id:%d\n", node_id, orcl_ctx_id);

    // 获取当前节点所关系的所有转发dct
    orcl_get_all_dc_target(orcl_ctx);

    
    // 创建一系列线程参数
    auto param_arr = new forward_thread_params[group_num];

    // 初始化线程参数
    for(i = 0; i < group_num; i++){
        if(orcl_ctx->group_id == i){
            // 自己不需要线程，也不需要线程参数
            continue;
        }

        param_arr[i].orcl_ctx = orcl_ctx;
        param_arr[i].tid = i;
    }
    
    // 为上下文一口气获取所有的远程qp的信息，
    // 分为两种，一种是normal的，一种是forward的
    get_all_published_qp(orcl_ctx);

    // 因为理论上队列只需要连接一次
    // 只有0号上下文需要执行一次连接
    // 和对应的节点建立连接
    // 只要pre_ctx是null就需要建立连接
    if(pre_ctx == NULL){
        orcl_all_qp_connect(orcl_ctx);
    }

    // 创建一系列线程，传入的是转发dct缓冲区，监听的缓冲区偏移，转发rc缓冲区，所有的转发RC
    // 不用转发自己组的
    // 当前要分配的核
    int numa_node_index = orcl_ctx->orcl_ctx_id * (group_num-1);

    for(i = 0; i < group_num; i++){
        if(orcl_ctx->group_id == i){
            continue;
        }

        // if(orcl_ctx->numa_node_id == 0){
        //     continue;
        // }

        // 创建线程，std::thread返回的是一个指针
        orcl_ctx->thread_arr[i] = std::thread(run_forward_thread, &param_arr[i]);

        // 这里创建线程亲和性，隔着处理器放置
        cpu_set_t cpuset;
        // 将cpuset全部置为0初始化
        CPU_ZERO(&cpuset);

        // uint64_t random = 0xbeef;

        CPU_SET(node_cpu_tp[numa_node_id][numa_node_index%NODE_CPU_COUNT], &cpuset);
        
        numa_node_index++;
        
        // 设置处理器亲和性，让一个线程优先运行在对应CPU_SET中
        int rc = pthread_setaffinity_np(orcl_ctx->thread_arr[i].native_handle(),
                                        sizeof(cpu_set_t), &cpuset);

        // 亲和性绑定
        if (rc != 0)
        {
            printf("Error %d while calling pthread_setaffinity_np\n", rc);
        }
    }

    // 这里设置一个全局同步，等待所有进程初始化完成
    char ready_str[ORCL_QP_NAME_SIZE] = {};

    sprintf(ready_str, "init_finish_%d", orcl_ctx->node_id);

    orcl_publish_ready(ready_str);

    // 等待所有的节点准备好
    for(i = 0; i < orcl_ctx->node_num; i++){
        if(i == orcl_ctx->node_id){
            continue;
        }

        char remote_ready_str[ORCL_QP_NAME_SIZE] = {};

        sprintf(remote_ready_str, "init_finish_%d", i);

        orcl_wait_till_ready(remote_ready_str);
    }

    return orcl_ctx;
}


void orcl_write_message_to_remote(struct orcl_context* orcl_ctx, int offset){

    // 首先查看内容对不对
    assert(orcl_ctx != NULL && offset % 512 == 0 && orcl_ctx->message_buffer != NULL);

    // 将头部读出来
    struct orcl_message_header* header = (struct orcl_message_header*)(orcl_ctx->message_buffer + offset);

    // 本地的缓冲区和远程的缓冲区都不能溢出，
    assert(offset + header->size < orcl_ctx->message_buffer_size && header->size < 512);
    
    assert(header->offset + header->size < orcl_ctx->message_buffer_size);

    assert(header->size > (int)sizeof(header));

    // 远程offset也要对齐，并且魔数也要正确
    assert(header->offset % 512 == 0 && header->magic == 'a' && header->node_id < orcl_ctx->node_num);

    // 根据路由策略，执行发送
    // 查看要发送的节点在哪一个组
    int remote_group_id = header->node_id / orcl_ctx->group_size;

    // 查看是不是当前组
    if(remote_group_id == orcl_ctx->group_id){
        // 当前组，直接使用normal rc来进行发送
        orcl_inner_group_rdma_write(orcl_ctx, header->node_id % orcl_ctx->group_size, offset, header->offset, header->size);
    }else{

        // 其他组，向其他组的转发节点发送
        // 首先查看要转发的位置
        orcl_dci_rdma_write(orcl_ctx, remote_group_id, offset, orcl_ctx->forward_offset[remote_group_id], header->size);
        // 这里转发的偏移量进行一次自增
        orcl_ctx->forward_offset[remote_group_id] = (orcl_ctx->forward_offset[remote_group_id] + 512) % FORWARD_DCT_BUF_SIZE;
    }
}

